import os
import pandas as pd


class DataOperation:
    # 传入文件路径,读取data数据
    def __init__(self, file_path, sheet=0):
        # 获取文件后缀判断文件类型
        _, extension = os.path.splitext(file_path)
        if extension in ('.txt', '.csv', '.tsv'):
            self.data = pd.read_csv(file_path)
        elif extension in ('.xls', '.xlsx'):
            self.data = pd.read_excel(file_path, sheet_name=sheet)
        # sql语句不知道怎么描述
        else:
            print('文件路径有误')

    def get_data_to_dict(self):
        # [{}, {}, {}, ...]
        return [self.data.loc[i].to_dict() for i in self.data.index.values]

    def get_data_tolist(self):
        # [[], [], [], ...]
        return self.data.values.tolist()

    def head(self, num=5):
        # 从data前面读取数据num行
        return self.data.head(num)

    def tail(self, num=5):
        # 从data尾部读取数据num行
        return self.data.tail(num)

    def write_data_to_excel(self, new_data, location=0):
        """
        new_data是要传入的新数据, 是列表格式
        location是传入的位置, 默认在最前面
        """
        self.data.loc[location] = new_data

    def shape(self):
        # 查看数据的形状,返回行,列
        return self.data.shape

    def columns(self):
        # 查看所有字段名称
        return self.data.columns

    def index(self):
        # 查看数据的索引
        return self.data.index


if __name__ == '__main__':
    data1 = DataOperation('data/冻结行列窗格_模板.xlsx', 1)
    # print(data1.head(3))
    # list1 = [0, '20220303', '0330测试班', 41, 41, '推荐']
    # data1.write_data_to_excel(list1)
    # print(data1.head())
    # print(data1.shape())
    # data1.columns()
    print(data1.index())
